import argparse

import numpy as np
import torch
import torch.nn as nn
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from torch.autograd import Variable
from torchvision.utils import save_image

import utils
from models import generator, discriminator

# 设置参数
parser = argparse.ArgumentParser()
parser.add_argument("--n_epochs", type=int, default=30, help="number of epochs of training")
parser.add_argument("--batch_size", type=int, default=128, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0001, help="learning rate")
parser.add_argument("--noise_dim", type=int, default=100, help="dimensionality of the noise")
parser.add_argument("--generator_path", type=str, default="./models/generator.pth", help="name of generator")
parser.add_argument("--discriminator_path", type=str, default="./models/discriminator.pth", help="name of discriminator")
parser.add_argument("--pic_path", type=str, default="./images", help="path of picture")
parser.add_argument("--normal", type=int, default=0, help="whether to use normal noise")
opt = parser.parse_args()
print(opt)

dataloader = utils.loadData(opt.batch_size)
D = discriminator().cuda()
G = generator(opt.noise_dim, 1 * 56 * 56).cuda()

criterion = nn.BCELoss()
d_optimizer = torch.optim.RMSprop(D.parameters(), lr=opt.lr)
g_optimizer = torch.optim.RMSprop(G.parameters(), lr=opt.lr)
Tensor = torch.cuda.FloatTensor
# # 可视化模型
# utils.draw_model(D, (opt.batch_size, 1, 28, 28))
# utils.draw_model(G, (opt.batch_size, opt.noise_dim))

# 训练模型
for epoch in range(opt.n_epochs):
    for i, (img, _) in enumerate(dataloader):
        img_num = img.size(0)
        # =================训练判别器
        real_img = Variable(img).cuda()
        real_label = Variable(torch.ones(img_num, 1)).cuda()
        fake_label = Variable(torch.zeros(img_num, 1)).cuda()

        # 计算真实图像的loss
        real_out = D(real_img)
        d_loss_real = criterion(real_out, real_label)

        # 计算“假”图像的loss
        if opt.normal == 0:
            noise = Variable(torch.randn(img_num, opt.noise_dim)).cuda()  # 随机分布
        else:
            noise = Variable(Tensor(np.random.normal(0, 1, (img_num, opt.noise_dim)))).cuda()  # 正态分布
        fake_img = G(noise)
        fake_out = D(fake_img)
        d_loss_fake = criterion(fake_out, fake_label)

        d_loss = d_loss_real + d_loss_fake
        d_optimizer.zero_grad()
        d_loss.backward()
        d_optimizer.step()

        # ===============训练生成器
        # 计算“假”图像的loss
        if opt.normal == 0:
            noise = Variable(torch.randn(img_num, opt.noise_dim)).cuda()  # 随机分布
        else:
            noise = Variable(Tensor(np.random.normal(0, 1, (img_num, opt.noise_dim)))).cuda()  # 正态分布
        fake_img = G(noise)
        output = D(fake_img)

        g_loss = criterion(output, real_label)
        g_optimizer.zero_grad()
        g_loss.backward()
        g_optimizer.step()

        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], d_loss: {:.6f}, g_loss: {:.6f} '
                  'D real: {:.6f}, D fake: {:.6f}'.format(epoch, opt.n_epochs, d_loss.item(), g_loss.item(),
                                                          real_out.data.mean(), fake_out.data.mean()))
    # 保存图片
    fake_images = utils.to_img(fake_img.cpu().data)
    save_image(fake_images, opt.pic_path + 'fake_images-{}.png'.format(epoch + 1))

# 保存模型
torch.save(G.state_dict(), opt.generator_path)
torch.save(D.state_dict(), opt.discriminator_path)
